import cv2
from yolov5 import YOLOv5

# 加载预训练的YOLOv5模型
model = YOLOv5("./yolov5s.pt",device='0')  # 选择模型

# 打开摄像头
cap = cv2.VideoCapture(0)

while True:
    # 从摄像头读取帧
    ret, frame = cap.read()

    if not ret:
        break

    # 使用YOLOv5进行目标检测
    results = model.predict(frame)

    # 在帧上绘制检测结果
    for *xyxy, conf, cls in results.xyxy[0]:
        label = f'{model.model.names[int(cls)]} {conf:.2f}'
        cv2.rectangle(frame, (int(xyxy[0]), int(xyxy[1])), (int(xyxy[2]), int(xyxy[3])), (0, 0, 255), 2)
        cv2.putText(frame, label, (int(xyxy[0]), int(xyxy[1]) - 10), cv2.FONT_HERSHEY_SIMPLEX, 1.0, (0, 0, 255), 2)

    # 显示帧


    # 按'q'键退出
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

# 释放资源并关闭窗口
cap.release()
cv2.destroyAllWindows()
